module ModuleAggregates
    
    use Globals
    use Toolbox
    
    contains

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!             Initialize Measure                   !!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    subroutine InitializeMeasure
    
    implicit none
    
    ! same share in each group
    mu = (1.0D0/dble(NS))
    mu = mu/sum(mu)
    
    end subroutine InitializeMeasure


    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!             Measure Computation  w/out Qmu       !!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    subroutine MeasureStationaryPara

    implicit none

    ! Local variable declarations
    real(8) :: munew(NS)
    real(8) :: a_np, a_np_vec(1), da
    integer :: aind_np, aind_np_vec(1), is, is_np, ix_np
    real(8) :: wwmu
    integer :: itmu, maxitmu, showmu
    integer:: ik,is_np_1,is_np_2,is_np_3
    
    ! Tolerance level, max number of iterations, slow updating parameter
    tolmu = 1e-11; maxitmu = 5000; wwmu = 0.00D0

    showmu = 1
    do itmu = 1, maxitmu
        
        munew = 0d0 ! initialize updated measure

        do is = 1, NS

            ! linear interpolation in the asset dimension
            a_np        = min(a_pol(is),amax)
            a_np_vec    = a_np
            !aind_np_vec = vsearch_FET(a_np_vec,a_vec)
            aind_np_vec = vsearchprm_FET(a_np_vec,a_vec,aprm)
            aind_np     = aind_np_vec(1)
            da = (a_np - a_vec(aind_np))/(a_vec(aind_np+1) - a_vec(aind_np))
            
            ! account for all possible future productivity states
            do ix_np = 1, Nx
                is_np = aind_np + Na*(ix_np-1)
                munew(is_np)   = munew(is_np)   + (1.0D0-da)*Px(xind(is),ix_np)*mu(is)
                munew(is_np+1) = munew(is_np+1) +    da     *Px(xind(is),ix_np)*mu(is)
            end do

        enddo
        
        ! Show if measure not equal to one
        if (abs(sum(munew)-1.0d0)>10D0*tolmu)  then
            print*, 'measure is unequal one', sum(munew)
        endif
          
        ! Error
        errmu = norm2(munew-mu)
        if (showmu==0) then
            !print *, 'itmu = ', itmu,', errmu = ', errmu
            showmu = 100        ! reset showmu
        end if

        ! Check convergence
        if (errmu <= tolmu) then
        !print *, 'Stationary measure converged at itmu = ', itmu
        exit
        else
        mu = wwmu*mu + (1-wwmu)*munew
        end if

        if (itmu == maxitmu) then
        print *, 'cannot find mu, ', errmu
        end if

        showmu = showmu - 1
    end do  ! end loop in itmu
    
    end subroutine MeasureStationaryPara

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!             Measure Aggregation         !!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  
    subroutine Aggregation
    
    implicit none
    
    ! Local variable declarations
    real(8), dimension(NS) :: yk, yl, ytot, transfers, et
    real(8) :: I0, I0_l, I0_k, I3, chit
    real(8), dimension(NS) ::  S1u, S2u
   
    S1u = S(:,1)        ! assets
    S2u = S(:,2)        ! productivity

    Kagg = max(sum(S1u*mu) - Debt,1d-8)     ! agg. capital; equal to sum(a_pol*mu) - D
    Lagg = sum(S2u*h_pol*mu)                ! agg. efficiency units of labor
    Hagg = sum(h_pol*mu)                    ! agg. labor
    Cagg = sum(c_pol*mu)                    ! agg. labor

    Yagg = (Kagg**(1.0D0-alpha))*(Lagg**alpha)                  ! agg. output
    r_new = - delta + (1.0D0-alpha)*((Lagg/Kagg)**(alpha))      ! implied interest rate
    
    yk  = r*S1u             ! capital income
    yl  = wge*S2u*h_pol     ! labor income

    ytot = yk+yl            ! total income

    ! Compute normalizing income for the tax function (normalization kept constant except for calibration)
    call mean_income(ytot,ymean_new)

    ! Tax revenue
    I0_l = sum(mu*exp(log(lambda)*((yl/ymean)**(-gamma)))*yl)   ! labor tax revenue
    I0_k = sum(mu*(tauk*yk))                                    ! capital tax revenue

    !  Transfer spending
    et   = exp(-rt*((yk+yl)/ymean-omegat))
    chit = exp(rt*omegat)/(1.0D0 + exp(rt*omegat))
    transfers = (mt*ymean)*(et/(1.0D0+et))*(1D0/chit)
    I3 = sum(mu*(transfers))

    ! Government budget constraint
    ! spending + interest expenses + transfers - capital tax revenue - labor tax revenue - consumption tax revenue
    BC  = (G + r*Debt + I3 - I0_k - I0_l - tauc*Cagg)

    ! Asset market clearing
    ! private saving - govt debt - capital demand
    asset = sum(S1u*mu) - Debt - ((r + delta)/(1.0D0-alpha))**(-1.0D0/alpha) * Lagg

    end subroutine Aggregation

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!             Mean Income                 !!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

     subroutine mean_income(y_in,mean_inc)

        implicit none

        real(8), intent(in)  :: y_in(NS)      ! total income vector input
        real(8), intent(out) :: mean_inc      ! output: measure of "mean" income
        real(8) :: mu_vec(NS), y_in_vec(NS)   ! for sorting
        real(8) :: mu_sorted(NS), y_in_sorted(NS), mu_cumu(NS)  ! for sorting
        integer :: ind_sort(NS)                                 ! for sorting
        integer :: p_50_ind, p_40_ind, p_60_ind, iter           ! percentile indices and counter

        mu_vec = mu
        y_in_vec = y_in

        ! Mean Income !
        if (norm_version == 1) then

            ! mean income
            mean_inc = sum(y_in_vec*mu_vec)
        
        ! Median Income !
        elseif (norm_version == 2) then

            ! median income
            ! sort total income array
            y_in_sorted = y_in_vec
            ind_sort = 0
            call hpsort_eps_epw(NS,y_in_sorted,ind_sort,1D-16)
            ! sort measure array correspondingly
            mu_sorted = mu_vec(ind_sort)
            ! compute cumulative measure
            mu_cumu = 0d0
            mu_cumu(1) = mu_sorted(1)
            do iter = 2, NS
                mu_cumu(iter) = mu_cumu(iter-1) + mu_sorted(iter)
            end do
            ! find index of the median and corresponding income value
            p_50_ind = minloc(abs(mu_cumu-0.5d0),dim=1)
            mean_inc = y_in_sorted(p_50_ind)

        ! Mean Income Third Quintile !
        elseif (norm_version == 3) then

            ! mean income of the third quintile
            ! sort total income array
            y_in_sorted = y_in_vec
            ind_sort = 0
            call hpsort_eps_epw(NS,y_in_sorted,ind_sort,1D-16)
            ! sort measure array correspondingly
            mu_sorted = mu_vec(ind_sort)
            ! compute cumulative measure
            mu_cumu = 0d0
            mu_cumu(1) = mu_sorted(1)
            do iter = 2, NS
                mu_cumu(iter) = mu_cumu(iter-1) + mu_sorted(iter)
            end do
            ! find indices of the 40th and 60th percentile
            p_40_ind = minloc(abs(mu_cumu-0.4d0),dim=1)
            p_60_ind = minloc(abs(mu_cumu-0.6d0),dim=1)
            ! compute mean income of the third quintile
            mean_inc = sum(y_in_sorted(p_40_ind:p_60_ind)*mu_sorted(p_40_ind:p_60_ind))/sum(mu_sorted(p_40_ind:p_60_ind))

        end if

    end subroutine mean_income
           
end module ModuleAggregates



